import numpy as np


class MinibatchKmeans:
    def __init__(self, n_clusters=8, batch_size=1024, max_iter=100, random_state=None):
        """
        初始化 MiniBatchKMeans 聚类器

        参数:
            n_clusters (int): 聚类的簇数量
            batch_size (int): 每次随机选取的小批量大小
            max_iter (int): 最大迭代次数
            random_state (int): 随机数种子，保证结果可复现
        """
        self.n_clusters = n_clusters
        self.batch_size = batch_size
        self.max_iter = max_iter
        self.random_state = random_state
        self.cluster_centers_ = None  # 用于存储聚类中心

    def fit(self, data):
        """
        使用 MiniBatchKMeans 对数据进行聚类

        参数:
            data (ndarray): 输入数据，形状为 (n_samples, n_features)
        """
        rng = np.random.RandomState(self.random_state)
        n_samples, n_features = data.shape

        # Step 1: 初始化簇中心
        self.cluster_centers_ = data[rng.choice(n_samples, self.n_clusters, replace=False)]

        # Step 2: 迭代优化
        for _ in range(self.max_iter):
            # 随机抽取小批量数据点
            minibatch_indices = rng.choice(n_samples, self.batch_size, replace=False)
            minibatch = data[minibatch_indices]

            # Step 3: 计算小批量数据点到簇中心的距离并分配标签
            distances = np.linalg.norm(minibatch[:, np.newaxis] - self.cluster_centers_, axis=2)
            labels = np.argmin(distances, axis=1)

            # Step 4: 更新簇中心
            for i in range(self.n_clusters):
                points_in_cluster = minibatch[labels == i]
                if len(points_in_cluster) > 0:  # 避免空簇
                    self.cluster_centers_[i] = points_in_cluster.mean(axis=0)

        return self

    def predict(self, data):
        """
        对每个数据点预测所属的簇

        参数:
            data (ndarray): 输入数据，形状为 (n_samples, n_features)

        返回:
            labels (ndarray): 每个数据点所属的簇标签，形状为 (n_samples,)
        """
        distances = np.linalg.norm(data[:, np.newaxis] - self.cluster_centers_, axis=2)
        return np.argmin(distances, axis=1)